import os
import sys
from typing import List, Tuple, Optional

import cv2
import numpy as np


def _rotate_bound(image: np.ndarray, angle_deg: float) -> np.ndarray:
	"""以不裁剪方式旋轉圖片。

	參考: 先計算旋轉後的邊界尺寸，再平移使其完整保留。
	"""
	(h, w) = image.shape[:2]
	center = (w / 2.0, h / 2.0)
	M = cv2.getRotationMatrix2D(center, angle_deg, 1.0)
	cos = abs(M[0, 0])
	sin = abs(M[0, 1])

	# 新尺寸
	nW = int((h * sin) + (w * cos))
	nH = int((h * cos) + (w * sin))

	# 調整平移，避免被裁剪
	M[0, 2] += (nW / 2) - center[0]
	M[1, 2] += (nH / 2) - center[1]
	return cv2.warpAffine(image, M, (nW, nH), flags=cv2.INTER_LINEAR, borderMode=cv2.BORDER_REPLICATE)


def _largest_rect_angle(gray: np.ndarray) -> Optional[float]:
	"""根據圖片中最大矩形框估計需要旋轉的角度（度）。

	返回值為應該順時針旋轉的角度（OpenCV 角度為順時針為正）。
	若無法檢出，返回 None。
	"""
	# 邊緣檢測
	blur = cv2.GaussianBlur(gray, (5, 5), 0)
	# 自適應二值化 + 邊緣，加強對細線條的魯棒性
	thr = cv2.adaptiveThreshold(blur, 255, cv2.ADAPTIVE_THRESH_GAUSSIAN_C, cv2.THRESH_BINARY_INV, 35, 15)
	edges = cv2.Canny(thr, 50, 150)

	# 連通、閉合
	kernel = cv2.getStructuringElement(cv2.MORPH_RECT, (5, 5))
	dil = cv2.dilate(edges, kernel, iterations=2)
	close = cv2.morphologyEx(dil, cv2.MORPH_CLOSE, kernel, iterations=2)

	contours, _ = cv2.findContours(close, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
	if not contours:
		return None

	h, w = gray.shape[:2]
	img_area = float(h * w)

	# 選擇最大且接近矩形的區域
	best = None
	best_area = 0.0
	for cnt in contours:
		area = cv2.contourArea(cnt)
		if area < img_area * 0.05:  # 忽略太小的
			continue
		rot = cv2.minAreaRect(cnt)
		(box_w, box_h) = rot[1]
		if box_w == 0 or box_h == 0:
			continue
		rect_area = box_w * box_h
		# 優先面積更大且更接近圖像尺寸的
		if rect_area > best_area:
			best = rot
			best_area = rect_area

	if best is None:
		return None

	# 解析角度：OpenCV 角度範圍 (-90, 0]
	angle = best[-1]
	box_w, box_h = best[1]
	# 將角度轉為「使長邊接近水平」的矯正角
	if box_w < box_h:
		rotation = angle  # 例如 -3.2 -> 逆時針 3.2，需要順時針旋轉為正值
	else:
		rotation = angle + 90

	# 傳回順時針為正（cv2.getRotationMatrix2D 即如此）
	return rotation


def correct_frame_skew(image: np.ndarray) -> np.ndarray:
	"""根據最大矩形圖框自動糾正傾斜，返回糾正後圖片。

	若無法檢測到矩形，將回退到基於投影的整體糾偏。
	"""
	if image is None:
		raise ValueError("image is None")
	gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY) if image.ndim == 3 else image.copy()

	angle = _largest_rect_angle(gray)
	if angle is None:
		# 回退方法：霍夫線估計整體文本傾斜
		edges = cv2.Canny(gray, 50, 150)
		lines = cv2.HoughLines(edges, 1, np.pi / 180.0, threshold=150)
		if lines is not None and len(lines) > 0:
			# 平均角度（相對於 x 軸的弧度）
			angles = []
			for rho_theta in lines[:50]:
				rho, theta = rho_theta[0]
				# 將角度轉為 -90~90 之間
				ang = (theta * 180.0 / np.pi) - 90
				if -45 <= ang <= 45:
					angles.append(ang)
			if angles:
				angle = float(np.median(angles))
			else:
				angle = 0.0
		else:
			angle = 0.0

	# OpenCV 正角為順時針旋轉；我們期望把歪斜角度抵消
	rotated = _rotate_bound(image, angle)
	return rotated


def _is_image_file(name: str) -> bool:
	name_l = name.lower()
	return name_l.endswith((".jpg", ".jpeg", ".png", ".bmp", ".tif", ".tiff", ".webp"))


def process_path(path: str) -> List[str]:
	"""處理單檔或資料夾，將糾正後圖片輸出到原目錄下的新子目錄 `corrected`。

	返回輸出文件列表。
	"""
	if not os.path.exists(path):
		raise FileNotFoundError(path)

	outputs: List[str] = []

	def process_file(img_path: str, out_dir: str) -> Optional[str]:
		img = cv2.imdecode(np.fromfile(img_path, dtype=np.uint8), cv2.IMREAD_COLOR)
		if img is None:
			return None
		res = correct_frame_skew(img)
		# 使用 imencode 以支援含中文路徑
		base = os.path.basename(img_path)
		name, ext = os.path.splitext(base)
		out_path = os.path.join(out_dir, f"{name}_corrected{ext if ext else '.png'}")
		success, buf = cv2.imencode(ext if ext else '.png', res)
		if not success:
			return None
		os.makedirs(out_dir, exist_ok=True)
		buf.tofile(out_path)
		return out_path

	if os.path.isfile(path):
		root_dir = os.path.dirname(path) or os.getcwd()
		out_dir = os.path.join(root_dir, "corrected")
		if _is_image_file(path):
			out = process_file(path, out_dir)
			if out:
				outputs.append(out)
	else:
		# 僅處理該資料夾下一層（不遞迴）
		out_dir = os.path.join(path, "corrected")
		for name in os.listdir(path):
			full = os.path.join(path, name)
			if os.path.isfile(full) and _is_image_file(name):
				out = process_file(full, out_dir)
				if out:
					outputs.append(out)

	return outputs


def _choose_path_interactive() -> Optional[str]:
	"""調用系統檔案對話方塊選擇單圖或資料夾。"""
	try:
		import tkinter as tk
		from tkinter import filedialog
		root = tk.Tk()
		root.withdraw()
		# 先問模式
		from tkinter import messagebox
		resp = messagebox.askyesno("選擇模式", "是：選擇文件夾\n否：選擇單張圖片")
		if resp:
			return filedialog.askdirectory(title="選擇圖片所在文件夾") or None
		else:
			return filedialog.askopenfilename(title="選擇單張圖片", filetypes=[("Images", "*.jpg;*.jpeg;*.png;*.bmp;*.tif;*.tiff;*.webp")]) or None
	except Exception:
		return None


def run_interactive() -> None:
	"""交互執行：選擇路徑並批量處理。"""
	path = _choose_path_interactive()
	if not path:
		print("未選擇任何文件或文件夾。")
		return
	outs = process_path(path)
	if outs:
		print(f"完成，輸出 {len(outs)} 個文件：")
		for p in outs:
			print(p)
	else:
		print("沒有處理到任何圖片。")


if __name__ == "__main__":
	if len(sys.argv) > 1:
		all_out = []
		for p in sys.argv[1:]:
			all_out.extend(process_path(p))
		print(f"完成，共輸出 {len(all_out)} 個文件。")
	else:
		run_interactive()


